import os
import cv2
import json
import numpy as np
import shutil
from tqdm.std import tqdm
def vis_parsing_maps(im, parsing_anno, save_im=False, save_path=''):
    stride=0
    # Colors for all 20 parts
    part_colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0],
                   [255, 0, 85], [255, 0, 170],
                   [0, 255, 0], [85, 255, 0], [170, 255, 0],
                   [0, 255, 85], [0, 255, 170],
                   [0, 0, 255], [85, 0, 255], [170, 0, 255],
                   [0, 85, 255], [0, 170, 255],
                   [255, 255, 0], [255, 255, 85], [255, 255, 170],
                   [255, 0, 255], [255, 85, 255], [255, 170, 255],
                   [0, 255, 255], [85, 255, 255], [170, 255, 255]]
    im = np.array(im)
    vis_im = im.copy().astype(np.uint8)
    vis_parsing_anno = parsing_anno.copy().astype(np.uint8)
    # vis_parsing_anno = cv2.resize(vis_parsing_anno, None, fx=stride, fy=stride, interpolation=cv2.INTER_NEAREST)
    vis_parsing_anno_color = np.zeros((vis_parsing_anno.shape[0], vis_parsing_anno.shape[1], 3)) + 255

    num_of_class = np.max(vis_parsing_anno)

    for pi in range(1, num_of_class + 1):
        index = np.where(vis_parsing_anno == pi)
        vis_parsing_anno_color[index[0], index[1], :] = part_colors[pi]

    vis_parsing_anno_color = vis_parsing_anno_color.astype(np.uint8)
    vis_im = cv2.addWeighted(cv2.cvtColor(vis_im, cv2.COLOR_RGB2BGR), 0.4, vis_parsing_anno_color, 0.6, 0)
    if save_im:
        cv2.imwrite(save_path, vis_im, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
data= open('./crop_test_face_new+13.json')
data = json.load(data)
print(data.keys())
print(len(data['images_name']))
for i, img_path in enumerate(tqdm(data['images_name'])):
    all_face_anno = img_path.replace('image', 'concat_img//anno').replace('jpg', 'png')
    part_face_anno = img_path.replace('test/image', 'test_crop_gaussian/ensemble/anno').replace('jpg', 'png')
    # part_hair_anno = img_path.replace('image', 'pred_hair_hat_model_transform_multi_scale/anno').replace('jpg', 'png')
    save_path = img_path.replace('image', 'concat_img/anno').replace('jpg', 'png')
    save_path_show = img_path.replace('image', 'concat_img/weight_img').replace('jpg', 'png')
    # print(save_path[:-13])
    if not os.path.exists(save_path[:-13]):
        os.mkdir(save_path[:-13])
    if not os.path.exists(save_path_show[:-13]):
        os.mkdir(save_path_show[:-13])
    boxes = data['predict_boxes'][i]
    all_anno = cv2.imread(all_face_anno, cv2.IMREAD_GRAYSCALE)
    part_anno = cv2.imread(part_face_anno, cv2.IMREAD_GRAYSCALE)
    # hair_anno = cv2.imread(part_hair_anno, cv2.IMREAD_GRAYSCALE)
    x=boxes['x']
    y=boxes['y']
    size=boxes['size']
    # print(part_face_anno, boxes, all_face_anno)
    #add for hair
    # all_anno[all_anno==10]=0
    # all_anno[all_anno==15]=0
    # all_anno[hair_anno==10]=10
    # all_anno[hair_anno==15]=15

    all_anno[y:y + size, x:x + size] = part_anno
    cv2.imwrite(save_path, all_anno)
    im_ori = cv2.imread(img_path)
    # vis_parsing_maps(im_ori, all_anno , save_im=True, save_path=save_path_show)

# ids = [os.path.join('data/cvpr/test/pred_image_GCnet/anno', path) for path in os.listdir('./data/cvpr/test/pred_image_GCnet/anno')]
# print('Done!')
# count = 0
# for n, id in enumerate(ids):
#     # print(id[-4:])
#     imgs = [os.path.join(id, path) for path in os.listdir(id)]
#     for i, im_path in enumerate(imgs):

#         # print(im_path,im_path.replace('pred_all_face_model_transform/anno', 'concat_img_transform/anno'))
#         if not os.path.exists(im_path.replace('pred_image_GCnet/anno', 'concat_img/anno')):
#             count += 1
#             print(im_path, im_path.replace('pred_image_GCnet/anno', 'concat_img/anno'))
#             shutil.copy(im_path, im_path.replace('pred_image_GCnet/anno', 'concat_img/anno'))
# print(count)

# for i, img in data[]:
# image = data['image']
# box = data['crop_boxes']
# root_dir = './data/cvpr/test/'
# data_dir = root_dir+'image/'
# ann_dir = root_dir+'/pred_crop_face_model/anno/'
# for id in os.listdir(ann_dir):
#     for img in os.listdir(data_dir+id):
#         im = cv2.imread(data_dir+id+'/'+img)
#         anno = cv2.imread(ann_dir+id+'/'+img.replace('jpg', 'png'), cv2.IMREAD_GRAYSCALE)
#         print(im,anno)
#         vis_parsing_maps(im, anno, save_im=True, save_path=root_dir+'vis', im_name=img)